from fastapi import WebSocket
from typing import List, Dict, Any, Set
import json
import asyncio
from datetime import datetime
from config.settings import DEFAULT_SYMBOL, DEFAULT_INST_TYPE, UPDATE_INTERVAL
from utils.logger import logger

class WebSocketManager:
    def __init__(self):
        self.active_connections = {}
        self.order_subscribers = {}
        self.ticker_subscribers = {}  # 新增：ticker订阅者字典
        self.is_running = True
        self.start_time: float = 0
        self.end_time: float = 0
        
    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections[id(websocket)] = websocket
        client_id = id(websocket)
        logger.info(f"新的WebSocket连接建立 - 客户端ID: {client_id}, 当前连接数: {len(self.active_connections)}")
        
    def disconnect(self, websocket: WebSocket):
        client_id = id(websocket)
        if client_id in self.order_subscribers:
            del self.order_subscribers[client_id]
            logger.info(f"已从订单订阅列表中移除客户端ID: {client_id}")
        if client_id in self.ticker_subscribers:
            del self.ticker_subscribers[client_id]
            logger.info(f"已从ticker订阅列表中移除客户端ID: {client_id}")
        logger.info(f"WebSocket连接断开 - 客户端ID: {client_id}, 当前连接数: {len(self.active_connections)}")
        
    def add_order_subscriber(self, websocket: WebSocket):
        """添加订单频道订阅者"""
        client_id = id(websocket)
        self.order_subscribers[client_id] = websocket
        logger.info(f"客户端订阅订单频道 - 客户端ID: {client_id}, 当前订阅数: {len(self.order_subscribers)}")
        
    def add_ticker_subscriber(self, websocket: WebSocket):
        """添加ticker频道订阅者"""
        client_id = id(websocket)
        self.ticker_subscribers[client_id] = websocket
        logger.info(f"客户端订阅ticker频道 - 客户端ID: {client_id}, 当前订阅数: {len(self.ticker_subscribers)}")
        
    def set_time_range(self, start_time: float, end_time: float):
        self.start_time = start_time
        self.end_time = end_time
        
    def start_simulation(self):
        """启动WebSocket模拟"""
        self.is_running = True
        logger.info(f"WebSocket模拟已启动，当前状态：")
        logger.info(f"- 运行状态: {self.is_running}")
        logger.info(f"- 活跃连接数: {len(self.active_connections)}")
        logger.info(f"- 订单订阅者数: {len(self.order_subscribers)}")
        logger.info(f"- Ticker订阅者数: {len(self.ticker_subscribers)}")
        
    def stop_simulation(self):
        """停止WebSocket模拟"""
        self.is_running = False
        self.order_subscribers.clear()
        self.ticker_subscribers.clear()
        logger.info("WebSocket模拟已停止")
        logger.info(f"- 活跃连接数: {len(self.active_connections)}")
        logger.info(f"- 订单订阅者数: {len(self.order_subscribers)}")
        logger.info(f"- Ticker订阅者数: {len(self.ticker_subscribers)}")
        
    async def push_order_update(self, order_type, order, price):
        """向订阅了orders频道的客户端推送订单更新"""
        # 组装消息体
        now = int(datetime.now().timestamp() * 1000)
        if order_type == "open":
            msg = {
                "action": "snapshot",
                "arg": {
                    "instType": "USDT-FUTURES",
                    "channel": "orders",
                    "instId": "default"
                },
                "data": [{
                    "orderId": order["order_id"],
                    "status": "filled",
                    "price": str(price),
                    "posSide": "long",
                    "tradeSide": "open",
                    "fillTime": now
                }],
                "ts": now
            }
        elif order_type == "close":
            msg = {
                "action": "snapshot",
                "arg": {
                    "instType": "USDT-FUTURES",
                    "channel": "orders",
                    "instId": "default"
                },
                "data": [{
                    "orderId": order["order_id"],
                    "clientOid": order["order_id"],
                    "status": "filled",
                    "price": str(price),
                    "posSide": "long",
                    "tradeSide": "close",
                    "fillTime": now
                }],
                "ts": now
            }
        
        # 记录订单更新信息
        subscriber_count = len(self.order_subscribers)
        logger.info(f"准备推送订单更新, 订单ID: {order['order_id']}, 类型: {order_type}, 价格: {price}, 订阅者数量: {subscriber_count}")
        
        if subscriber_count == 0:
            logger.warning("没有找到订阅者，无法推送订单更新")
            return
            
        # 记录订阅者客户端ID
        client_ids = list(self.order_subscribers.keys())
        logger.info(f"当前订阅者客户端ID列表: {client_ids}")
        
        # 推送给所有订阅了订单频道的客户端
        disconnected_clients = []
        for client_id in client_ids:
            try:
                if client_id not in self.order_subscribers:
                    logger.warning(f"客户端ID: {client_id}不在订阅列表中，跳过")
                    continue
                    
                ws = self.order_subscribers[client_id]
                await ws.send_json(msg)
                logger.info(f"成功推送订单更新到客户端ID: {client_id}")
            except KeyError:
                logger.warning(f"未找到客户端ID: {client_id}，可能已断开连接")
                disconnected_clients.append(client_id)
            except Exception as e:
                logger.error(f"推送订单更新失败 - 客户端ID: {client_id}, 错误: {str(e)}")
                disconnected_clients.append(client_id)
        
        # 清理已断开连接的客户端
        for client_id in disconnected_clients:
            if client_id in self.order_subscribers:
                del self.order_subscribers[client_id]
                logger.info(f"从订阅列表中移除断开连接的客户端: {client_id}, 剩余订阅者: {len(self.order_subscribers)}")
        
    async def push_ticker_update(self, price: float):
        """向订阅了ticker频道的客户端推送价格更新"""
        # 组装ticker消息
        now = int(datetime.now().timestamp() * 1000)
        msg = {
            "action": "snapshot",
            "arg": {
                "instType": "USDT-FUTURES",
                "channel": "ticker",
                "instId": "ETHUSDT"
            },
            "data": [{
                "instId": "ETHUSDT",
                "last": f"{price:.2f}",
                "lastSz": "0.01",
                "askPx": f"{price + 0.5:.2f}",
                "askSz": "1.0",
                "bidPx": f"{price - 0.5:.2f}",
                "bidSz": "1.0",
                "volume24h": "1000.0",
                "ts": str(now)
            }]
        }
        
        # 记录推送信息
        subscriber_count = len(self.ticker_subscribers)
        logger.info(f"准备推送Ticker更新, 价格: {price}, 订阅者数量: {subscriber_count}")
        
        if subscriber_count == 0:
            logger.warning("没有找到Ticker订阅者，无法推送价格更新")
            return
            
        # 记录订阅者客户端ID
        client_ids = list(self.ticker_subscribers.keys())
        logger.info(f"当前ticker订阅者客户端ID列表: {client_ids}")
        
        # 向所有订阅者推送
        disconnected_clients = []
        for client_id in client_ids:
            try:
                if client_id not in self.ticker_subscribers:
                    logger.warning(f"客户端ID: {client_id}不在ticker订阅列表中，跳过")
                    continue
                    
                ws = self.ticker_subscribers[client_id]
                await ws.send_json(msg)
                logger.info(f"成功推送Ticker更新到客户端ID: {client_id}, 价格: {price}")
            except KeyError:
                logger.warning(f"未找到客户端ID: {client_id}，可能已断开连接")
                disconnected_clients.append(client_id)
            except Exception as e:
                logger.error(f"推送Ticker更新失败 - 客户端ID: {client_id}, 错误: {str(e)}")
                disconnected_clients.append(client_id)
        
        # 清理已断开连接的客户端
        for client_id in disconnected_clients:
            if client_id in self.ticker_subscribers:
                del self.ticker_subscribers[client_id]
                logger.info(f"从ticker订阅列表中移除断开连接的客户端: {client_id}, 剩余订阅者: {len(self.ticker_subscribers)}")
                
    async def broadcast(self, message: Dict[str, Any]):
        """广播消息给所有连接的客户端"""
        message_str = json.dumps(message, ensure_ascii=False)
        logger.info(f"[WebSocket广播] 广播消息内容: {message_str[:200]}...")
        
        for connection in self.active_connections.values():
            try:
                await connection.send_json(message)
                client_id = id(connection)
                logger.info(f"[WebSocket回应] 客户端ID: {client_id}, 响应内容: {message_str[:200]}...")
            except Exception as e:
                client_id = id(connection)
                logger.error(f"广播消息发送失败 - 客户端ID: {client_id}, 错误: {str(e)}")
                
    def get_connection_count(self) -> int:
        """获取当前WebSocket连接数量"""
        return len(self.active_connections)
        
    async def handle_ping(self, websocket: WebSocket):
        """处理ping请求"""
        client_id = id(websocket)
        logger.info(f"[WebSocket回应] 客户端ID: {client_id}, ping -> pong")
        await websocket.send_text("pong")
        
    async def handle_login(self, websocket: WebSocket):
        """处理登录请求"""
        client_id = id(websocket)
        response = {
            "event": "login",
            "code": "0",
            "msg": "登录成功"
        }
        logger.info(f"[WebSocket回应] 客户端ID: {client_id}, 响应内容: {json.dumps(response, ensure_ascii=False)}")
        await websocket.send_json(response)

    async def handle_message(self, websocket: WebSocket, message: str):
        """处理接收到的WebSocket消息"""
        client_id = id(websocket)
        logger.info(f"[WebSocket接收] 客户端ID: {client_id}, 消息内容: {message}")
        
        try:
            # 如果消息是字符串"ping"，直接返回"pong"
            if message == "ping":
                logger.info(f"[WebSocket接收] 客户端ID: {client_id}, 收到ping请求")
                await self.handle_ping(websocket)
                return
                
            data = json.loads(message)
            
            if data.get('op') == 'ping':
                logger.info(f"[WebSocket接收] 客户端ID: {client_id}, 收到ping请求")
                await self.handle_ping(websocket)
            elif data.get('op') == 'login':
                await self.handle_login(websocket)
            else:
                # 处理其他类型的消息
                response = {
                    "status": "success",
                    "message": "received",
                    "requestData": data
                }
                logger.info(f"[WebSocket回应] 客户端ID: {client_id}, 响应内容: {json.dumps(response, ensure_ascii=False)}")
                await websocket.send_json(response)
            
        except json.JSONDecodeError as e:
            error_response = {
                "status": "error",
                "message": "无效的JSON格式",
                "error": str(e)
            }
            logger.error(f"[WebSocket错误] 客户端ID: {client_id}, JSON解析错误: {str(e)}")
            await websocket.send_json(error_response)
            
        except Exception as e:
            error_response = {
                "status": "error",
                "message": "处理请求时发生错误",
                "error": str(e)
            }
            logger.error(f"[WebSocket错误] 客户端ID: {client_id}, 处理错误: {str(e)}")
            await websocket.send_json(error_response) 